{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# DiscreteDistribution"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "license_cell",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\nAtlanta, Georgia 30332-0415\nAll Rights Reserved\n\nAuthors: Frank Dellaert, et al. (see THANKS for the full author list)\n\nSee LICENSE for the license information"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/discrete/doc/DiscreteDistribution.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `DiscreteDistribution` is a `DiscreteConditional` with no parents. It represents a prior probability on a single discrete variable, e.g., $P(X)$. It can be thought of as a special case of a factor, specifically one that sums to 1.\n",
        "\n",
        "It derives from `gtsam.DiscreteConditional` and can be used interchangeably with a `DecisionTreeFactor` or `DiscreteConditional` in most contexts, such as adding it to a `DiscreteFactorGraph` or `DiscreteBayesNet`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "import graphviz\n",
        "\n",
        "from gtsam.symbol_shorthand import X"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating a DiscreteDistribution\n",
        "\n",
        "A `DiscreteDistribution` is defined on a single `DiscreteKey`, which is a tuple of a gtsam `Key` (integer) and the variable's cardinality (number of possible outcomes). They can be constructed in several ways."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Created from string:\n",
            "Discrete Prior\n",
            " P( x0 ):\n",
            " Choice(x0) \n",
            " 0 Leaf  0.1\n",
            " 1 Leaf  0.9\n",
            "\n",
            "\n",
            "Created from vector:\n",
            "Discrete Prior\n",
            " P( x0 ):\n",
            " Choice(x0) \n",
            " 0 Leaf  0.1\n",
            " 1 Leaf  0.9\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Define a key for a binary variable (cardinality=2)\n",
        "X0 = (X(0), 2)\n",
        "\n",
        "# --- Method 1: From a spec string ---\n",
        "# The string \"1/9\" specifies the probabilities for X0=0 and X0=1.\n",
        "dd_string = gtsam.DiscreteDistribution(X0, \"1/9\")\n",
        "print(\"Created from string:\")\n",
        "dd_string.print()\n",
        "\n",
        "# --- Method 2: From a list of probabilities ---\n",
        "dd_vector = gtsam.DiscreteDistribution(X0, [0.1, 0.9])\n",
        "print(\"\\nCreated from vector:\")\n",
        "dd_vector.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Operations on DiscreteDistribution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "P(X0=0) = 0.1\n",
            "P(X0=1) = 0.9\n",
            "\n",
            "PMF vector: [0.1, 0.9]\n"
          ]
        }
      ],
      "source": [
        "# --- Evaluate the probability of a specific outcome ---\n",
        "# This uses the __call__ method.\n",
        "prob_X0_is_0 = dd_string(0)\n",
        "prob_X0_is_1 = dd_string(1)\n",
        "print(f\"P(X0=0) = {prob_X0_is_0}\")\n",
        "print(f\"P(X0=1) = {prob_X0_is_1}\")\n",
        "\n",
        "# --- Retrieve the Probability Mass Function (PMF) ---\n",
        "pmf = dd_string.pmf()\n",
        "print(f\"\\nPMF vector: {pmf}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Since `DiscreteDistribution` is a type of factor, it can be visualized."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Title: G Pages: 1 -->\n",
              "<svg width=\"134pt\" height=\"109pt\"\n",
              " viewBox=\"0.00 0.00 134.00 109.46\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 105.46)\">\n",
              "<title>G</title>\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-105.46 130,-105.46 130,4 -4,4\"/>\n",
              "<!-- 0x1247da968 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>0x1247da968</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"63\" cy=\"-81.33\" rx=\"20.13\" ry=\"20.13\"/>\n",
              "<text text-anchor=\"middle\" x=\"63\" y=\"-76.28\" font-family=\"Times,serif\" font-size=\"14.00\">x0</text>\n",
              "</g>\n",
              "<!-- 0x1247e7e50 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>0x1247e7e50</title>\n",
              "<polygon fill=\"none\" stroke=\"black\" points=\"54,-25.2 0,-25.2 0,0 54,0 54,-25.2\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-7.55\" font-family=\"Times,serif\" font-size=\"14.00\">0.10</text>\n",
              "</g>\n",
              "<!-- 0x1247da968&#45;&gt;0x1247e7e50 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>0x1247da968&#45;&gt;0x1247e7e50</title>\n",
              "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M53.73,-63.15C49.23,-54.8 43.76,-44.67 38.95,-35.75\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"42.06,-34.15 34.23,-27.01 35.9,-37.47 42.06,-34.15\"/>\n",
              "</g>\n",
              "<!-- 0x1247ab110 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>0x1247ab110</title>\n",
              "<polygon fill=\"none\" stroke=\"black\" points=\"126,-25.2 72,-25.2 72,0 126,0 126,-25.2\"/>\n",
              "<text text-anchor=\"middle\" x=\"99\" y=\"-7.55\" font-family=\"Times,serif\" font-size=\"14.00\">0.90</text>\n",
              "</g>\n",
              "<!-- 0x1247da968&#45;&gt;0x1247ab110 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>0x1247da968&#45;&gt;0x1247ab110</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M72.27,-63.15C76.77,-54.8 82.24,-44.67 87.05,-35.75\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"90.1,-37.47 91.77,-27.01 83.94,-34.15 90.1,-37.47\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x124c509b0>"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Visualize the factor\n",
        "# For a single-variable factor, this is just a single node.\n",
        "graphviz.Source(dd_string.dot())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Displaying with Markdown/HTML\n",
        "In a Jupyter environment, GTSAM objects often have a rich HTML display."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<p>  <i>P(x0):</i></p>\n",
              "<div>\n",
              "<table class='DecisionTreeFactor'>\n",
              "  <thead>\n",
              "    <tr><th>x0</th><th>value</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>0</th><td>0.1</td></tr>\n",
              "    <tr><th>1</th><td>0.9</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/markdown": [
              " *P(x0):*\n",
              "\n",
              "|x0|value|\n",
              "|:-:|:-:|\n",
              "|0|0.1|\n",
              "|1|0.9|\n"
            ],
            "text/plain": [
              "Discrete Prior\n",
              " P( x0 ):\n",
              " Choice(x0) \n",
              " 0 Leaf  0.1\n",
              " 1 Leaf  0.9\n"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dd_string"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}